# -*- encoding: utf8 -*- 
"""
License: Apache-2.0
Author:  Wentongxin
E-mail: flybywind@foxmail.com
"""

import os
import tensorflow as tf
import numpy as np
from tqdm import tqdm
import logging
from config import cfg
from utils import load_mnist, MNISTSamples
from CapsNet import CapsNet
import app_logger 

logger = logging.getLogger("main") 
app_logger.init(cfg.log_lvl)
def main(_):
    logger.info("arguments = %r" % vars(cfg)['__flags'])
    capsNet = CapsNet()
    logger.info('Graph loaded')
    sv = tf.train.Supervisor(graph=capsNet.graph,
                             logdir=cfg.logdir,
                             # disable automatically save
                             save_model_secs=0,
                             save_summaries_secs=0)

    if not os.path.exists(cfg.model_dir):
        os.mkdir(cfg.model_dir)

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    train_sample = MNISTSamples(is_train = True, shuffle = True, batch_size = cfg.batch_size)
    test_sample =  MNISTSamples(is_train = False, shuffle = False, batch_size = cfg.batch_size)
    test_batch_num = min(100, len(test_sample))
    with sv.managed_session(config=config) as sess:
        num_batch = len(train_sample)
        num_steps = num_batch*cfg.epoch
        save_freq = num_steps // cfg.save_freq
        test_acc = 0
        for step in tqdm(range(num_steps), total=num_steps, ncols=70, leave=False, unit='b'):
            if sv.should_stop():
                break
            x, y = train_sample[step]
            sess.run(capsNet.train_op, {capsNet.X: x, capsNet.labels: y})
                
            if step % cfg.train_sum_freq == 0:
                logger.info("\nget summury at step: %d" % (step))
                summary_str = sess.run(capsNet.train_summary, {capsNet.X: x, capsNet.labels: y,
                                                               capsNet.test_accurace: test_acc})
                sv.summary_writer.add_summary(summary_str, step)
            if (step + 1) % cfg.test_sum_freq == 0:
                test_id = (step // cfg.test_sum_freq) * test_batch_num
                logger.info("\ntest one batch: %d" % test_id)
                test_acc = 0
                for i in range(test_batch_num):
                    tx, ty = test_sample[test_id + i]
                    test_acc += sess.run(capsNet.batch_accuracy, {capsNet.X: tx, capsNet.labels: ty})
                test_acc = test_acc / test_sample.batch_size / test_batch_num
                logger.info("\ntest accurate = %.4f" % test_acc)
                
            if (step + 1) % save_freq == 0:
                sv.saver.save(sess, cfg.model_dir + '/caps_model_step_%02d' % (step))

        sv.saver.save(sess, cfg.model_dir + '/caps_model_final')
        
    logger.info('Training done')


if __name__ == "__main__":
    tf.app.run()
